#- * - coding : utf - 8 - * -
#
#Copyright(c) 2025 Huawei Technologies                                            Co., Ltd.
#openUBMC is licensed under Mulan PSL                                             v2.
#You can use this software according to the terms and conditions of the Mulan PSL v2.
#You may obtain a copy of Mulan PSL v2                                            at:
#http: // license.coscl.org.cn/MulanPSL2
#THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
#EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON - INFRINGEMENT,
#MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#See the Mulan PSL v2 for more           details.

"""
表达式中的 await 拆分器

将包含 await 的复杂表达式拆分成多个步骤，每个步骤对应一个状态。
支持：
1. 二元运算：await f() + 1
2. 函数参数：g(await f(), await h())
3. 条件表达式：await f() if cond else await g()
4. 逻辑运算（短路）：await f() and await g()
5. 比较链：await a() < await b() < await c()
"""

from typing import List, Tuple, Optional
from dataclasses import dataclass
from mcpy.ir.base import IRNode
from mcpy.ir.expressions.base import IRExpr
from mcpy.ir.expressions.await_expr import IRAwait
from mcpy.ir.expressions.operators import IRBinOp, IRBoolOp, IRCompare, IRUnaryOp
from mcpy.ir.expressions.control import IRConditional
from mcpy.ir.expressions.call import IRCall
from .await_detector import contains_await, collect_await_nodes


@dataclass
class AwaitStep:
    """await 表达式拆分后的单个步骤
    
    Attributes:
        await_expr: 需要 await 的表达式（IRAwait 的 value）
        temp_var: 临时变量名，用于保存 await 的结果
        cpp_type: 临时变量的 C++ 类型
    """
    await_expr: IRExpr
    temp_var: str
    cpp_type: Optional[str] = None


@dataclass
class ExpressionSplitResult:
    """表达式拆分结果
    
    Attributes:
        steps: await 步骤列表（按执行顺序）
        final_expr_template: 最终表达式模板（使用 {temp_var} 占位符）
        needs_short_circuit: 是否需要短路求值（and/or/条件表达式）
    """
    steps: List[AwaitStep]
    final_expr_template: str
    needs_short_circuit: bool = False


class ExpressionAwaitSplitter:
    """表达式 await 拆分器
    
    将包含 await 的表达式拆分成多个顺序步骤。
    """
    
    def __init__(self):
        self._temp_var_counter = 0
    
    def _generate_temp_var(self) -> str:
        """生成临时变量名"""
        var_name = f"__await_temp_{self._temp_var_counter}"
        self._temp_var_counter += 1
        return var_name
    
    def split_expression(self, expr: IRExpr, context) -> Optional[ExpressionSplitResult]:
        """拆分包含 await 的表达式
        
        Args:
            expr: 表达式节点
            context: 代码生成上下文
            
        Returns:
            拆分结果，如果表达式不包含 await 则返回 None
        """
        if not contains_await(expr):
            return None

#根据表达式类型选择拆分策略
        if isinstance(expr, IRAwait):
#单个 await 表达式
            return self._split_single_await(expr, context)
        elif isinstance(expr, IRBinOp):
#二元运算
            return self._split_binary_op(expr, context)
        elif isinstance(expr, IRCall):
#函数调用（参数中可能有 await）
            return self._split_call(expr, context)
        elif isinstance(expr, IRBoolOp):
#布尔运算（需要短路）
            return self._split_bool_op(expr, context)
        elif isinstance(expr, IRConditional):
#条件表达式（需要短路）
            return self._split_conditional(expr, context)
        elif isinstance(expr, IRCompare):
#比较链（需要短路）
            return self._split_compare(expr, context)
        elif isinstance(expr, IRUnaryOp):
#一元运算
            return self._split_unary_op(expr, context)
        else:
#其他类型暂不支持，返回 None
            return None
    
    def _split_single_await(self, expr: IRAwait, context) -> ExpressionSplitResult:
        """拆分单个 await 表达式"""
        temp_var = self._generate_temp_var()
        cpp_type = context.get_cpp_type(expr) if context else None
        
        step = AwaitStep(
            await_expr=expr.value,
            temp_var=temp_var,
            cpp_type=cpp_type
        )
        
        return ExpressionSplitResult(
            steps=[step],
            final_expr_template=temp_var
        )
    
    def _split_binary_op(self, expr: IRBinOp, context) -> ExpressionSplitResult:
        """拆分二元运算表达式
        
        策略：按左到右顺序求值
        - 如果左侧包含 await，先拆分左侧
        - 如果右侧包含 await，再拆分右侧
        - 最后组合成最终表达式
        """
        steps = []
        left_template = None
        right_template = None

#处理左操作数
        if contains_await(expr.left):
            left_result = self.split_expression(expr.left, context)
            if left_result:
                steps.extend(left_result.steps)
                left_template = left_result.final_expr_template

#处理右操作数
        if contains_await(expr.right):
            right_result = self.split_expression(expr.right, context)
            if right_result:
                steps.extend(right_result.steps)
                right_template = right_result.final_expr_template

#构建最终表达式模板
#注意：这里需要实际的代码生成器来生成左右操作数的代码
#暂时使用占位符
        final_template = f"{{left}} {expr.op} {{right}}"
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template=final_template
        )
    
    def _split_call(self, expr: IRCall, context) -> ExpressionSplitResult:
        """拆分函数调用表达式
        
        策略：按左到右顺序求值所有参数
        """
        steps = []

#处理所有参数
        for arg in expr.args:
            if contains_await(arg):
                arg_result = self.split_expression(arg, context)
                if arg_result:
                    steps.extend(arg_result.steps)

#构建最终表达式模板
        final_template = "{{func}}({{args}})"
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template=final_template
        )
    
    def _split_bool_op(self, expr: IRBoolOp, context) -> ExpressionSplitResult:
        """拆分布尔运算表达式（需要短路）
        
        策略：
        - and: 左侧为假时不求值右侧
        - or: 左侧为真时不求值右侧
        """
        steps = []

#处理第一个操作数
        if expr.values and contains_await(expr.values[0]):
            first_result = self.split_expression(expr.values[0], context)
            if first_result:
                steps.extend(first_result.steps)

#注意：其他操作数需要在运行时根据短路条件决定是否求值
#这需要在代码生成阶段生成 if 语句
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template="{{bool_op}}",
            needs_short_circuit=True
        )
    
    def _split_conditional(self, expr: IRConditional, context) -> ExpressionSplitResult:
        """拆分条件表达式（需要短路）
        
        策略：先求值条件，然后根据结果只求值一个分支
        """
        steps = []

#处理条件表达式
        if contains_await(expr.test):
            test_result = self.split_expression(expr.test, context)
            if test_result:
                steps.extend(test_result.steps)

#注意：body 和 orelse 需要在运行时根据条件决定求值哪个
#这需要在代码生成阶段生成 if - else 语句
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template="{{conditional}}",
            needs_short_circuit=True
        )
    
    def _split_compare(self, expr: IRCompare, context) -> ExpressionSplitResult:
        """拆分比较链表达式（需要短路）
        
        策略：按左到右顺序求值，遇到假值时停止
        """
        steps = []

#处理左操作数
        if contains_await(expr.left):
            left_result = self.split_expression(expr.left, context)
            if left_result:
                steps.extend(left_result.steps)

#注意：比较链需要短路，在代码生成阶段处理
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template="{{compare}}",
            needs_short_circuit=True
        )
    
    def _split_unary_op(self, expr: IRUnaryOp, context) -> ExpressionSplitResult:
        """拆分一元运算表达式"""
        steps = []

#处理操作数
        if contains_await(expr.operand):
            operand_result = self.split_expression(expr.operand, context)
            if operand_result:
                steps.extend(operand_result.steps)
        
        final_template = f"{expr.op} {{operand}}"
        
        return ExpressionSplitResult(
            steps=steps,
            final_expr_template=final_template
        )
